import pandas as pd
import xgboost as xgb
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from scipy.stats import spearmanr

def get_data():
    data = pd.read_csv('../../data/raw/train.csv')
    data.info()
    data.BusinessTravel = data.BusinessTravel.map({'Travel_Rarely': 1, 'Travel_Frequently': 2, 'Non-Travel': 0})
    data.MaritalStatus = data.MaritalStatus.map({'Divorced': 2, 'Single': 0, 'Married': 1})
    onehot = pd.get_dummies(data[['Department', 'EducationField', 'JobRole', 'OverTime']])
    data = pd.concat([data, onehot], axis=1)
    data = data.drop(columns=['Gender', 'Department', 'EducationField', 'JobRole', 'OverTime', 'Over18','OverTime_No'])
    # data.info()   43
    data['RoleStagnationPeriod'] = data['YearsInCurrentRole']/(data['YearsAtCompany'] + 10**-5) # 保留
    data['ManagerPeriod'] = data['YearsWithCurrManager']/(data['YearsAtCompany'] + 10**-5)      # 保留
    data['IncomeCompetitionRate'] = data['MonthlyIncome']/(data['TotalWorkingYears'] + 10**-5)  # 保留
    data['AgeWorkingRate'] = data['Age']/(data['TotalWorkingYears'] + 10**-5)      # 保留
    data['']

    data.info()
    # data_AgeWorkingRate = data.groupby('AgeWorkingRate', as_index=False)['Attrition'].mean()
    # spearmanr(data_AgeWorkingRate)
    # rho, p_value = spearmanr(data_AgeWorkingRate.AgeWorkingRate, data_AgeWorkingRate.Attrition)
    # print(rho, p_value)

    # data.to_csv('../../data/raw/train_ini.csv', index=False)
    return data


def xgbfeather(data):
    es = xgb.XGBClassifier(objective='binary:logistic')
    x = data.iloc[:,1:]
    y = data.iloc[:,0]
    es.fit(x, y)
    importances = es.feature_importances_
    df = pd.DataFrame({'Feature': data.columns[1:], 'Importance': importances})
    df.sort_values('Importance', ascending=False, inplace=True)
    print(df)


if __name__ == '__main__':
    data = get_data()
    # xgbfeather(data)
